from dataclasses import dataclass
from datetime import datetime
from requests import request
from uuid import uuid4
from json import dumps
from .response_models import (ModelsResponseModel, 
                              TokensCountResponseModel,
                              CompletionResponseModel,
                              CreateEmbeddingResponseModel)
from .request_models import (TokensCountRequestModel, 
                             CompletionsRequestModel,
                             CreateEmbeddingRequestModel)


class GigaChat:
    
    @dataclass
    class AuthToken():
        value: str
        expiration_date: datetime

    def __init__(self, credentials, scope):
        self.token = self.AuthToken(
            value = None,
            expiration_date = None
        )
        self.credentials = credentials
        self.scope = scope
        
    def authorize(func):
        def wrapper(self, *args, **kwargs):
            is_authorized: bool = lambda: self.token.value != None and self.token.expiration_date > datetime.now()
            auth_url: str = 'https://ngw.devices.sberbank.ru:9443/api/v2/oauth'
            if not is_authorized():
                payload={
                    'scope': self.scope,
                }
                headers = {
                    'Content-Type': 'application/x-www-form-urlencoded',
                    'Accept': 'application/json',
                    'RqUID': str(uuid4()),
                    'Authorization': f'Basic {self.credentials}'
                }
                response = request("POST", auth_url, headers=headers, data=payload)
                if response.status_code == 200:
                    self.token.value = response.json()['access_token']
                    self.token.expiration_date = datetime.fromtimestamp(response.json()['expires_at']/1000)
            return func(self, *args, **kwargs)
        return wrapper
    
    #Получение доступных моделей
    @authorize
    def get_models(self) -> ModelsResponseModel:
        url: str = 'https://gigachat.devices.sberbank.ru/api/v1/models'
        headers = {
            'Accept': 'application/json',
            'Authorization': f'Bearer {self.token.value}'
        }
        response = request("GET", url, headers=headers)
        if response.status_code == 200:
            payload = response.json()
            return ModelsResponseModel(
                object=payload['object'],
                data=[ModelsResponseModel.Data(
                    id=model['id'], 
                    object=model['object'], 
                    owned_by=model['owned_by']) for model in payload['data']]
            )
               

    #Получение количества токенов по входным сообщениям
    @authorize
    def get_tokens_count(self, input_model: TokensCountRequestModel) -> TokensCountResponseModel:
        url: str = 'https://gigachat.devices.sberbank.ru/api/v1/tokens/count'
        payload = dumps({
             "model": input_model.model,
             "input": input_model.input
        })
        headers = {
        'Content-Type': 'application/json',
        'Accept': 'application/json',
        'Authorization': f'Bearer {self.token.value}'
        }
        response = request("POST", url, headers=headers, data=payload)
        if response.status_code == 200:
            return [TokensCountResponseModel(tokens=tk['tokens'], characters=tk['characters']) for tk in response.json()]
        
 

    #Получение ответа модели
    @authorize
    def get_completions(self, input_model: CompletionsRequestModel) -> CompletionResponseModel:
        url: str = 'https://gigachat.devices.sberbank.ru/api/v1/chat/completions'
        payload = dumps({
            "model": input_model.model,
            "messages": input_model.messages,
            "temperature": input_model.temperature,
            "top_p": input_model.top_p,
            "n": input_model.n,
            "stream": input_model.stream,
            "max_tokens": input_model.max_tokens,
            "repetition_penalty": input_model.repetition_penalty,
            "update_interval": input_model.update_interval,
            }
        )
        headers = {
        'Content-Type': 'application/json',
        'Accept': 'application/json',
        'Authorization': f'Bearer {self.token.value}'
        }
        response = request("POST", url, headers=headers, data=payload)
        if response.status_code == 200:
            payload = response.json()
            return CompletionResponseModel(
                choices=[
                    CompletionResponseModel.Choice(
                        index=choice['index'],
                        finish_reason=choice['finish_reason'],
                        message=CompletionResponseModel.Choice.Message(
                            role=choice['message']['role'],
                            content=choice['message']['content'],
                            data_for_context=choice['message'].get('data_for_context'),
                        ),
                    ) for choice in payload['choices']
                ],
                created=payload['created'],
                model=payload['model'],
                object=payload['object'],
                usage=CompletionResponseModel.Usage(
                    prompt_tokens=payload['usage']['prompt_tokens'],
                    completion_tokens=payload['usage']['completion_tokens'],
                    total_tokens=payload['usage']['total_tokens'],
                    system_tokens=payload['usage'].get('system_tokens'),
                )
            )
        return response.json()
    
    def create_embeddings(self, request_model: CreateEmbeddingRequestModel) -> CreateEmbeddingResponseModel:
        url: str = 'https://gigachat.devices.sberbank.ru/api/v1/chat/embeddings'
        payload = dumps({
            'model': request_model.model,
            'input': request_model.input,
        })
        headers = {
            'Content-Type': 'application/json',
            'Accept': 'application/json',
            'Authorization': 'Bearer <TOKEN>'
            }
        response = request("POST", url, headers=headers, data=payload)
        if response.status_code == 200:
            content = response.json()
            return CreateEmbeddingResponseModel(
                object=content['object'],
                model=content['model'],
                data=[CreateEmbeddingResponseModel.Data( 
                    object=data_object['object'],
                    embedding=data_object['embedding'],
                    index=data_object['index'],
                    usage=CreateEmbeddingResponseModel.Data.Usage(prompt_tokens=data_object['usage']['prompt_tokens']),
                    ) for data_object in content['data']]
            )
        if response.status_code == 401:
        

    def get_image(self):
        pass

        